#_*_ coding:utf-8 _*_
#2018年4月13日，本程序完全自己所写
import numpy as np
import matplotlib.pyplot as plt
import os
import time
import math
from openpyxl import Workbook#用来编辑Excel
from sklearn.decomposition import PCA#添加包含PCA的库

titlefont={'family': 'Times New Roman',#设置字体格式
           'color':  'black',
           'weight': 'normal',
           'size': 12,
        }
xylabelfont={'family': 'Times New Roman',
             'color': 'black',
             'weight': 'normal',
             'size': 12
        }
labelfont={'family': 'Times New Roman',
           'size': 12
        }
realKey=[0x65,0xB1]#真实的轮密钥
startTime=time.clock()
class fileOper:
    '''获得当前目录下的功耗痕迹文件和明文文件名，并返回'''
    def getFileName(self):
        '''获得当前目录下功耗痕迹文件名和明文文件名，返回值：traceFile,textinFile'''
        curDir=os.getcwd()
        dataDir=[dir for dir in os.listdir(curDir) if os.path.isdir(dir) if dir.endswith('data')]

        if np.size(dataDir)!=1:
            print('请在当前目录存放一个数据文件')
            exit(1)
        else:
            dataDir=curDir+'\\'+dataDir[0]+'\\'+'traces'

            fileNameTemple=[filename for filename in os.listdir(dataDir) if filename.endswith ('traces.npy')]
            traceFile=dataDir+'\\'+fileNameTemple[0]

            fileNameTemple=[filename for filename in os.listdir(dataDir) if filename.endswith ('textin.npy')]
            textinFile=dataDir+'\\'+fileNameTemple[0]
            return traceFile,textinFile
    
class CPA:
    '''读取能耗和明文数据，并进行CPA攻击'''
    def __init__(self,traceIn,plaintextIn,):
        self.trace=traceIn#功耗
        self.plaintext=plaintextIn#明文
        self.traceNum=np.shape(self.trace)[0]#行数
        self.pointNum=np.shape(self.trace)[1]#列数
        self.sbox = (
            0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76,
            0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0,
            0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15,
            0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a, 0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75,
            0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84,
            0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b, 0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf,
            0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85, 0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8,
            0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5, 0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2,
            0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73,
            0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88, 0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb,
            0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79,
            0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08,
            0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a,
            0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e,
            0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf,
            0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16)
        self.HW = [bin(n).count("1") for n in range(0, 256)]  # 产生计算‘1’个数的查询表
        self.hypPower=np.zeros((self.traceNum,256),dtype='int8')#针对第subKey个子密钥而生成的猜测能耗
        self.CPAResult=np.zeros((256,32))#用来存储最终的攻击结果，
        self.keySelected=0#观察该密钥的破解结果

        self.pltUsed=[i for i in range(5,self.traceNum,10)]#计算相关系数时所用的明文组数[10,15,... 295]
        self.corrVsPoint=np.zeros((256,self.pointNum))#每种可能密钥在每个采样点对应的相关性
        self.corrVsPltNum=np.zeros((256,np.shape(self.pltUsed)[0]))#每种可能密钥在不同的明文组数时所取得的相关系数

    def calcCorrcoef(self,x, y=None, rowvar=True):
        '''计算相关系数'''
        c = np.cov(x, y,rowvar)
        try:
            d = np.diag(c)
        except ValueError:
            # scalar covariance
            # nan if incorrect value (nan, inf, 0), 1 otherwise
            return c / c
        stddev = np.sqrt(d)
    
        if (0 in stddev):
            c=np.array([[0,0],[0,0]])
        else:
            c /= stddev[:, None]
            c /= stddev[None, :]
        return c
    
    def calcIntermediate(self,pt, keyGuess):  # 产生中间值
        return self.sbox[pt ^ keyGuess]  # AES逐位异或并且通过s-box转换的中间值模型

    #计算针对第subKey个密钥字节的假设能耗Dx256
    def calcHypPower(self,subKey):
        for traceIndex in range(self.traceNum):
            for keyGuess in range(256):#对密钥的256种可能
                self.hypPower[traceIndex][keyGuess]=self.HW[self.calcIntermediate(self.plaintext[traceIndex][subKey],keyGuess)]

    #进行CPA攻击
    def attack(self,funSel='a',pltNum=100,stratPoint=0,endPoint=5000):#攻击停止字节（因为攻击16个字节的密钥太耗费时间），和所用明文组数，funsel选择功能 0只进行攻击，1相关系数随采样点之间的关系，2相关性随明文组数之间的关系
        oneByteResult=np.zeros((256,3))#存储对于subKey这个字节的攻击结果，第1列代表密钥的可能取值，第2列代表对应密钥取得最大相关系数的采样点，第3列代表所对应的相关系数
        correlationTemp=np.zeros((endPoint))#用于存储keyGuess和能耗矩阵中每一列的相关系数
        if funSel=='a':#function 0
            for subKey in range(16):#分别对每个选中的字节进行攻击
                self.calcHypPower(subKey)#计算假设功耗
                for keyGuess in range(256):#计算相关系数矩阵
                    for pointIndex in range(stratPoint,endPoint):
                        #correlationTemp[pointIndex]=np.absolute(self.calcCorrcoef(self.trace[:,pointIndex],self.hypPower[:,keyGuess])[0][1])
                        correlationTemp[pointIndex]=np.absolute(self.calcCorrcoef(self.trace[0:int(pltNum),pointIndex],self.hypPower[0:int(pltNum),keyGuess])[0][1])

                    oneByteResult[keyGuess,0]=keyGuess#可能值
                    oneByteResult[keyGuess,1]=np.argmax(correlationTemp)#取得最大值的索引
                    loc=np.argmax(correlationTemp)
                    #oneByteResult[keyGuess,2]=self.hypPower[keyGuess,temp[keyGuess,1]]#相关系数
                    oneByteResult[keyGuess,2]=correlationTemp[loc]#相关系数
                    
                index=np.argsort(oneByteResult[:,2])#对相关系数进行排序
                index=index[::-1]#由大到小排列
                for i in range(256):
                    self.CPAResult[i,subKey*2]=oneByteResult[index[i],0]
                    self.CPAResult[i,subKey*2+1]=oneByteResult[index[i],2]
               
        elif funSel=='b':#function 1,2
            subKey=self.keySelected
            self.calcHypPower(subKey)#计算假设功耗
            for keyGuess in range(256):#计算相关系数矩阵
                for pointIndex in range(stratPoint,endPoint):
                    #correlationTemp[pointIndex]=np.absolute(self.calcCorrcoef(self.trace[:,pointIndex],self.hypPower[:,keyGuess])[0][1])
                    correlationTemp[pointIndex]=np.absolute(self.calcCorrcoef(self.trace[0:int(pltNum),pointIndex],self.hypPower[0:int(pltNum),keyGuess])[0][1])

                oneByteResult[keyGuess,0]=keyGuess#可能值
                oneByteResult[keyGuess,1]=np.argmax(correlationTemp)#取得最大值的索引
                loc=np.argmax(correlationTemp)
                #oneByteResult[keyGuess,2]=self.hypPower[keyGuess,temp[keyGuess,1]]#相关系数
                oneByteResult[keyGuess,2]=correlationTemp[loc]#相关系数
                #figure1
                self.corrVsPoint[keyGuess,:]=correlationTemp[:]
        elif funSel=='c':
            #figure2
            subKey=self.keySelected
            self.calcHypPower(subKey)#计算假设功耗
            for keyGuess in range(256):#计算相关系数矩阵
                for pointIndex in range(stratPoint,endPoint):
                    #correlationTemp[pointIndex]=np.absolute(self.calcCorrcoef(self.trace[:,pointIndex],self.hypPower[:,keyGuess])[0][1])
                    correlationTemp[pointIndex]=np.absolute(self.calcCorrcoef(self.trace[0:int(pltNum),pointIndex],self.hypPower[0:int(pltNum),keyGuess])[0][1])

                oneByteResult[keyGuess,0]=keyGuess#可能值
                oneByteResult[keyGuess,1]=np.argmax(correlationTemp)#取得最大值的索引
                loc=np.argmax(correlationTemp)
                #oneByteResult[keyGuess,2]=self.hypPower[keyGuess,temp[keyGuess,1]]#相关系数
                oneByteResult[keyGuess,2]=correlationTemp[loc]#相关系数
            self.corrVsPltNum[:,self.jj]=oneByteResult[:,2]

    #将攻击结果存入Excel
    def saveResult(self):
        wb=Workbook()
        ws=wb.active
        ws.title="攻击结果"
        for subKey in range(13):#Excel 中的A-Z列
            ws.merge_cells("%c1:%c1"%(65+2*subKey,65+2*subKey+1))
            ws['%c1'%(65+2*subKey)]=subKey+1

            for i in range(256):
                ws['%c%d'%(65+2*subKey,2+i)]='0x%02X'%int(self.CPAResult[i][subKey*2])
                ws['%c%d'%(65+2*subKey+1,2+i)]=self.CPAResult[i][subKey*2+1]
        for subKey in range(3):
            ws.merge_cells("A%c1:A%c1"%(65+2*subKey,65+2*subKey+1))
            ws['A%c1'%(65+2*subKey)]=14+subKey

            for i in range(256):
                ws['A%c%d'%(65+2*subKey,2+i)]='0x%02X'%int(self.CPAResult[i][26+subKey*2])
                ws['A%c%d'%(65+2*subKey+1,2+i)]=self.CPAResult[i][26+subKey*2+1]
        wb.save('Helloworld.xlsx')
        
    #绘制出加密一条明文所产生的能耗
    def plotTrace(self):
        '''绘制出加密一条明文所产生的能耗'''
        fig3=plt.figure(3,figsize=(4.5,3))
        ax3=fig3.add_axes([0.12,0.15,0.8,0.75])
        ax3.plot(self.trace[0,:])
        plt.xlabel("Sampling point",fontdict=xylabelfont)
        plt.ylabel("Power consumption",fontdict=xylabelfont)
        plt.xlim(0,self.pointNum)
        ax3.tick_params(labelsize=9)
        filename='PowerConsumption.jpg'
        plt.savefig(filename,dpi=400)

    #绘制出相关系数同采样点之间的关系
    def plotCorrelationVsPoint(self):
        '''绘制出相关系数同采样点之间的关系'''
        self.attack('b',100)
        fig1=plt.figure(1,figsize=(4.5,3))
        ax1=fig1.add_axes([0.12,0.15,0.8,0.75])
        x=range(0,self.pointNum)
        for i in range(0,256):
            ax1.plot(x,self.corrVsPoint[i,:],color='xkcd:light blue')
        actualSelectedKey=realKey[self.keySelected]#所选密钥字节的真实值
        ax1.plot(x,self.corrVsPoint[actualSelectedKey,:])
        #plt.plot(x,self.corrVsPoint[actualSelectedKey,:])
        plt.xlabel("Sampling point",fontdict=xylabelfont)
        plt.ylabel("Correlation coefficient",fontdict=xylabelfont)
        plt.ylim(0,1)
        plt.xlim(0,self.pointNum)
        ax1.tick_params(labelsize=9)
        #plt.title("Correlation ",fontdict=titlefont)
        filename='CorrelationVsPoint.jpg'
        plt.savefig(filename,dpi=400)
        #plt.show()

    def plotCorrelationVsPlaintext(self):
        '''绘制出相关系数同明文组数之间的关系，本实验很浪费时间'''
        self.jj=0
        for i in range(np.shape(self.pltUsed)[0]):
            self.attack('c',self.pltUsed[i])
            self.jj=self.jj+1
        fig2=plt.figure(2,figsize=(4.5,3))
        ax2=fig2.add_axes([0.12,0.15,0.8,0.75])
        for i in range(0,256):
            ax2.plot(self.pltUsed,self.corrVsPltNum[i,:],color='xkcd:light blue')
        actualSelectedKey=realKey[self.keySelected]#所选密钥字节的真实值
        ax2.plot(self.pltUsed,self.corrVsPltNum[actualSelectedKey,:])

        plt.xlabel("Plaintext number ",fontdict=xylabelfont)
        plt.ylabel("Correlation coefficient",fontdict=xylabelfont)
        plt.ylim(0,1)
        plt.xlim(0,self.traceNum)
        ax2.tick_params(labelsize=9)
        #plt.title("Correlation ",fontdict=titlefont)
        filename='CorrelationVsPltNum.jpg'
        plt.savefig(filename,dpi=400)

class preTrans:
    def calcPCA(self,trace,componNum=100):
        '''进行主成分分析，输入参数依次为要进行PCA的矩阵，和要缩减到的维数'''
        pca=PCA(n_components=componNum)
        newTrace=pca.fit_transform(trace)
        print('贡献率为：%f'%(np.sum(pca.explained_variance_ratio_)))
        return newTrace
#测试模块
fileOpen=fileOper()
traceFile,textinFile=fileOpen.getFileName()
trace=np.load(traceFile)#加载功耗痕迹
plaintext=np.load(textinFile)#加载明文
#prework=preTrans()#创建功耗痕迹预处理对象
#newTrace=prework.calcPCA(trace,50)#进行主成分分析法降维
#print(np.shape(newTrace))

cpa=CPA(trace,plaintext)
cpa.plotTrace()
cpa.plotCorrelationVsPoint()
cpa.plotCorrelationVsPlaintext()
#cpa.attack(0,100,16)
#cpa.saveResult()
#cpa.plotCorrelation()
endTime=time.clock()
print('Runtime is %f seconds' %(endTime-startTime))
plt.show()